from pathlib import Path
from pyrsistent import m, pmap
from toolz import unique
import os
from pydicom import multival, DataElement
import numpy as np
import pydicom
from collections import defaultdict
private_tags = m(
    bvalue={
        "GE MEDICAL SYSTEMS": (0x0043, 0x1039),
        "SIEMENS": (0x0019, 0x000C),
        "TOSHIBA": (0x0018, 0x9087),
        "PHILIPS MEDICAL SYSTEMS": (0x2001, 0x1003),
        "NEUSOFT MEDICAL SYSTEMS CO., LTD.": (0x5501, 0x1231),
    }
)


class DCMHelper:
    def get_dcms(self, path, files=[]):
        for x in path.iterdir():
            if x.is_dir():
                self.get_dcms(x, files)
            elif x.suffix == '.dcm':
                files.append(x.resolve())
        return files

    def get_dcm_tags(self, image_files: list):

        if not image_files:
            print('Input files is empty!')
            return []

        tag_list = []
        for f in image_files:
            if not os.path.exists(f):
                print(f'{f} not exists, will skip!')
                continue

            try:
                ds = pydicom.dcmread(f, force=True, stop_before_pixels=True)
                assert ds.get('Rows') and ds.get('Columns'), \
                    f'Dcm must has "Rows" and "Columns": {f}'
            except Exception as e:
                print(e)
                continue

            metadata_json = {}
            for key in ds.dir():
                try:
                    metadata_json[key] = ds.get(key, '')
                except Exception as e:
                    print(e)

            metadata_json['FilePath'] = f
            metadata_json['0019100A'] = ds.get_item('0019100A')
            # metadata_json['AnatomicalPlane'] = utils.anatomical_plane_ds(ds)
            bvalue_tag = self.try_to_get_bvalue(ds)
            metadata_json['bvalue'] = bvalue_tag

            tag_list.append(metadata_json)
        return tag_list

    def try_to_get_bvalue(self, ds):
        manufacturer = str(ds.get('Manufacturer')).strip().upper()
        bvalue_tag = private_tags.bvalue.get(manufacturer)

        if bvalue_tag:
            element = ds.get(bvalue_tag, None)
            if isinstance(element, DataElement):
                val = element.value

                try:
                    if isinstance(val, float):
                        bvalue = val
                    elif isinstance(val, multival.MultiValue):
                        bvalue = val[0]
                    elif isinstance(val, list):
                        bvalue = val[0]
                    elif isinstance(val, bytes):
                        bvalue = val.decode('utf-8')
                    else:
                        bvalue = val

                    if isinstance(val, str):
                        bvalue = bvalue.split('\\', 1)[0]

                    return int(float(bvalue))
                except BaseException:
                    pass

        return -1


class SeriesSpliter:
    
    def split_multi_time_series(self, multi_time_series):
        """
        序列拆分->volume级别
        """
        if len(multi_time_series) < 1:
            return []

        elif len(multi_time_series) <= 2:
            return [[x for x in multi_time_series]]

        # only DTI series has the value '0019100a' in the dicom tags
        # return one volume directly if is a DTI series
        if multi_time_series[0].get("0019100A"):
            return [[x for x in multi_time_series]]

        # init
        first_position = multi_time_series[0].get("ImagePositionPatient", [0, 0, 0])
        bvalues = {
            y: i + 1 for i, y in enumerate(set([x.get('bvalue') for x in multi_time_series]))
        }
        delta = 0.1
        count = 0
        maxcount = 0
        group_time_series = defaultdict(list)

        # modes - order of priority (TriggerTime, TemporalPositionIdentifier, bvalue, ImagePositionPatient)
        mode = None
        for info in multi_time_series:
            if info.get('TriggerTime'):
                mode = "TriggerTime"
                break
            if (
                    info.get("TemporalPositionIdentifier", None) is not None
                    and info.get("NumberOfTemporalPositions") > 1
            ):
                mode = "TemporalPositionIdentifier"
                break
        if mode is None:
            if len(bvalues) > 1:
                mode = "bvalue"
            else:
                mode = "ImagePositionPatient"
        print(f"using mode: {mode}")
        # add first 2 instances
        # prev_acquisition_number = multi_time_series[1].get("AcquisitionNumber")
        prev_trigger_time = multi_time_series[1].get("TriggerTime")
        prev_trigger_time_trend = self.calculate_trend(
            multi_time_series[1].get("TriggerTime"), multi_time_series[0].get("TriggerTime")
        )

        if mode == "bvalue":
            for x in multi_time_series[:2]:
                count = bvalues[x.get("bvalue")]
                maxcount = count if count > maxcount else maxcount
                group_time_series["vol<{0}".format(count)].append(x)
        else:
            if mode == "TemporalPositionIdentifier":
                for y in multi_time_series[:2]:
                    count = y.get("TemporalPositionIdentifier", 0)
                    group_time_series["vol<{0}".format(count)].append(y)

            else:
                count += 1
                group_time_series["vol<{0}".format(count)].append(multi_time_series[0])
                group_time_series["vol<{0}".format(count)].append(multi_time_series[1])
            maxcount = count

        # only split series if there is repeated slice location
        # if the same location have two trigger time then need to split
        skip_split_volume = False
        if mode == "TriggerTime":
            # need to improve this to have buffer of 0.001 for slice location
            all_slice_location = list(unique(map(lambda s: s.get("SliceLocation", 0), multi_time_series)))
            if len(all_slice_location) == len(multi_time_series):
                skip_split_volume = True
                group_time_series["vol<{0}".format(count)] += multi_time_series[2:]

        if not skip_split_volume:
            # 2nd onwards
            for info in multi_time_series[2:]:
                pos = info.get("ImagePositionPatient", [0, 0, 0])
                # file_name = info.get("FilePath")
                # acquisition_number = info.get("AcquisitionNumber")
                bvalue = info.get("bvalue")
                temporal_position = info.get("TemporalPositionIdentifier")
                # num_temporal_position = info.get("NumberOfTemporalPositions")
                trigger_time = info.get("TriggerTime")
                trigger_time_trend = self.calculate_trend(
                    trigger_time, prev_trigger_time)

                # use trigger time to determine if the cycle repeated
                if (
                        ((prev_trigger_time_trend == "eq" and trigger_time_trend == "inc")
                         or (prev_trigger_time_trend == "inc" and trigger_time_trend == "dec"))
                        and mode == "TriggerTime"
                ):
                    count += 1
                    maxcount = count if count > maxcount else maxcount

                # use temporal position to determine if the cycle repeated
                elif temporal_position is not None and mode == "TemporalPositionIdentifier":
                    count = temporal_position
                    maxcount = count if count > maxcount else maxcount
                # use bvalue to determine if cycle repeated
                elif bvalue is not None and mode == "bvalue":
                    count = bvalues[bvalue]
                    maxcount = count if count > maxcount else maxcount

                # use image position to determine if the cycle repeated
                elif (
                        abs(pos[0] - first_position[0])
                        + abs(pos[1] - first_position[1])
                        + abs(pos[2] - first_position[2])
                        < delta
                ) and mode == "ImagePositionPatient":
                    count += 1
                    maxcount = count if count > maxcount else maxcount

                # set prevs
                # prev_acquisition_number = acquisition_number
                prev_trigger_time = trigger_time
                prev_trigger_time_trend = trigger_time_trend

                # add
                group_time_series["vol<{0}".format(count)].append(info)
        # 检查否有不平行数据
        self.check_volume_is_parallel(group_time_series)
        # add total volume
        # return list(group_time_series.values())

        # sort_by_positions
        return [self.sort_by_positions(vals) for vals in group_time_series.values()]

    def check_volume_is_parallel(self, group_time_series):
        for vol_num, dcm_infos in group_time_series.items():
            subgroup_timeseries = defaultdict(list)
            first_imageOrientation = list(map(float, dcm_infos[0].get("ImageOrientationPatient", [0, 0, 0, 0, 0, 0])))
            first_norm_vector = np.cross(first_imageOrientation[:3], first_imageOrientation[3:])
            subgroup_timeseries[0] = [first_norm_vector, [dcm_infos[0]]]  # 根据是否平行划分组别
            for dcm_info in dcm_infos[1:]:
                flag = False
                for v in subgroup_timeseries.values():
                    imageOrientation = list(map(float, dcm_info.get("ImageOrientationPatient", [0, 0, 0, 0, 0, 0])))
                    norm_vector = np.cross(imageOrientation[:3], imageOrientation[3:])
                    cosangle = np.dot(v[0], norm_vector) / \
                        (np.linalg.norm(v[0]) * np.linalg.norm(norm_vector))
                    if cosangle - np.cos(np.pi * 5 / 180) > 0:  # 小于5度（cos(5)=0.996）的默认为平行
                        v[1].append(dcm_info)
                        flag = True
                if not flag:
                    new_index = max(subgroup_timeseries.keys()) + 1
                    subgroup_timeseries[new_index] = [norm_vector, [dcm_info]]
            # 默认dicom数多的为正确序列, 初始为零
            max_num_dcms = 0
            for k, v in subgroup_timeseries.items():
                if len(v[-1]) >= max_num_dcms:
                    group_time_series[vol_num] = v[-1]
                    max_num_dcms = len(v[-1])

            time_series_len = len(group_time_series[vol_num])
            if time_series_len < len(dcm_infos):
                series_info = f"series_uid: {dcm_infos[0]['SeriesInstanceUID']}, study_uid: {dcm_infos[0]['StudyInstanceUID']}"
                print(
                    f"{group_time_series[vol_num][0].get('SOPInstanceUID')}_{time_series_len} have {len(dcm_infos) - time_series_len} non-parallel data, " +
                    series_info)

    def calculate_trend(self, cur, prev):
        if cur is None and prev is not None:
            return "dec"
        elif cur is not None and prev is None:
            return "inc"
        elif cur is None and prev is None:
            return "eq"
        elif cur < prev:
            return "dec"
        elif cur > prev:
            return "inc"
        else:
            return "eq"

    def sort_by_positions(self, image_tags):
        return list(
            sorted(
                image_tags,
                key=lambda x: self.transform_orientation(
                    x.get("ImagePositionPatient", [0, 0, 0]),
                    x.get("ImageOrientationPatient", [0, 0, 0, 0, 0, 0])
                )
            )
        )

    def transform_orientation(self, image_position, image_orientation):
        """
        get a sorted value according to dicom ImageOrientationPatient and PatientOrientation
        image_position: [-176.1845703125, -288.6845703125, -512.5]
        image_orientation: [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]
        """
        image_orientation = [float(x) for x in image_orientation]
        image_position = [float(x) for x in image_position]

        orient = np.array(image_orientation).reshape(2, 3)
        norm = np.cross(orient[0], orient[1])
        if np.dot(norm, [0, 0, 1]) < 0 and abs(np.dot(norm, [0, 0, 1])) > 0.707:
            norm = [-x for x in norm]
        return float(np.dot(image_position, norm))

if __name__ == '__main__':
    dcm_path = '/lfs/biomind/orthanc/original/1.2.840.2013.1030.20210829095159.210829000106/1.3.46.670589.11.33936.5.0.6196.2021082910114137033'
    dcm_path = Path(dcm_path)
    # 获取dcm tags
    helper = DCMHelper()
    file_list = helper.get_dcms(dcm_path)
    img_tags_list = helper.get_dcm_tags(file_list)
    # 拆分序列
    group_time_series = SeriesSpliter().split_multi_time_series(img_tags_list)

    print("dcm count: ", len(file_list))
    for i, volumes in enumerate(group_time_series, start=1):
        print(f"volume {i}, bvalue: {set([volume['bvalue'] for volume in volumes])}")
        print(f"file list: {set([volume['FilePath'].name for volume in volumes])}")
        print("\n")
